#!/usr/bin/python

from ceorModule import *

######################################
################ MAIN ################
## Is a 2-D box with lenght xL yL
if __name__ == "__main__":
  # And God says...
  N  = 2
  KT_eq = 0.1
  eps = 1.0
  sigma = 1.0
  delta_r_ini = 2.5*sigma
  dt = 0.001
  gas = Gas2D(10, 1, N, KT_eq, delta_r_ini)
  
  #--- Plot initial positions
  gas.arrayPos()
  gas.arrayVel()
  
  fig01 = plt.figure(1)
  plt.grid()
  plt.xlabel('$ x $')
  plt.ylabel('$ y $')
  plt.scatter(gas.X, gas.Y, marker='+')
  plt.xlim((0,gas.xlen+1))
  plt.ylim((0,gas.ylen+1))
  outputfilename="Fig01-XY_ini.png"
  plt.savefig(outputfilename)
  
  #--- Plot initial velocity distribution
  fig02 = plt.figure(2)
  gas.arraySpeed()
  #plt.hist(gas.VX, alpha=0.5)
  #plt.hist(gas.VY, alpha=0.5)
  plt.hist(gas.speed, alpha=0.5)
  outputfilename="Fig02-HistVxVySpeed_ini.png"
  plt.savefig(outputfilename)
  
  # 1. Calculamos F(t)
  gas.calculateForces(eps, sigma)
  print gas.Ri2
  print gas.Ri12 -gas.Ri6
  print 24*eps*((sigma**12)*gas.Ri12 -(sigma**6)*gas.Ri6)
  print 24*eps*((sigma**12)*gas.Ri12 -(sigma**6)*gas.Ri6)[0][1].sum()
  gas.calculateEkin()
  print '-'*20
  print 'E_tot =', gas.Ekin+gas.Epot , 'Ekin_ini = ', gas.Ekin, 'Epot_ini = ', gas.Epot
  
  # ------ Begin SIMULATION ------
  total_time = 50000#4900
  ttime_print = 1
  ttime = range(total_time)
  timeDim = total_time #/ttime_print + 1
  xt  = np.zeros([gas.N, timeDim])
  yt  = np.zeros([gas.N, timeDim])
  vxt = np.zeros([gas.N, timeDim])
  vyt = np.zeros([gas.N, timeDim])
  fxt = np.zeros([gas.N, timeDim])
  fyt = np.zeros([gas.N, timeDim])
  Ekint = np.zeros(timeDim)
  Epott = np.zeros(timeDim)
  Etott = np.zeros(timeDim)
  #######str(n).zfill(6)

  for i in ttime:
    gas.save_old() # fx_old=fx_new
    gas.stepPosition(dt) # input old variables = > output x_new
    gas.inBox()
    gas.calculateForces(eps, sigma)
    gas.stepVelocity(dt)
    # Save every each ttime_print
    #if i % ttime_print == 0 :
    i_print = i #/ttime_print
    gas.calculateEkin()
    Ekint[i_print] = gas.Ekin
    Epott[i_print] = gas.Epot
    #save particle trajectories
    for j in range(gas.N):
      xt[j][i_print]  = gas.particle[j].x_new
      yt[j][i_print]  = gas.particle[j].y_new
      vxt[j][i_print] = gas.particle[j].vx_new
      vyt[j][i_print] = gas.particle[j].vy_new
      fxt[j][i_print] = gas.particle[j].fx_new
      fyt[j][i_print] = gas.particle[j].fy_new
  # ------ End SIMULATION ------
  ttime = dt*np.array(ttime)
  Etott = Ekint+Epott
  gas.calculateEkin()
  print '-'*20
  print 'E_tot =', gas.Ekin+gas.Epot , 'Ekin_end = ', gas.Ekin, 'Epot_end = ', gas.Epot
  print '-- deltas --'*3
  print eps, dt, total_time, Etott.mean(), Etott.max() - Etott.min(),Etott.max(), Etott.min()
  #--- Plot Final positions
  gas.arrayPos()
  gas.arrayVel()
  fig03 = plt.figure(3)
  plt.grid()
  plt.xlabel('$ x $')
  plt.ylabel('$ y $')
  plt.scatter(gas.X, gas.Y, marker='o')
  plt.xlim((0,gas.xlen+1))
  plt.ylim((0,gas.ylen+1))
  outputfilename="Fig03-XY_end.png"
  plt.savefig(outputfilename)
  
  #--- Plot initial velocity distribution
  fig04 = plt.figure(4)
  plt.grid()
  gas.arraySpeed()
  #plt.hist(gas.VX, alpha=0.5)
  #plt.hist(gas.VY, alpha=0.5)
  plt.hist(gas.speed, alpha=0.5)
  outputfilename="Fig04-HistVxVySpeed_end.png"
  plt.savefig(outputfilename)
  
  
  fig05 = plt.figure(5)
  plt.grid()
  plt.xlabel('$ t $')
  plt.ylabel('$ E_{kin} $')
  plt.plot(ttime, Ekint,marker='.')
  outputfilename="Fig05-Ekint.png"
  plt.savefig(outputfilename)
  
  fig06 = plt.figure(6)
  plt.grid()
  plt.xlabel('$ t $')
  #plt.ylabel('$ E_{pot} $')
  plt.plot(ttime, -Epott)
  plt.plot(ttime,  Ekint)
  #plt.plot(ttime, (0.0615234375)*np.ones(total_time))
  #plt.plot(ttime, (1.0-0.0615234375)*np.ones(total_time))
  #plt.legend(["$-E_{pot}$","$E_{kin}$"])
  plt.xlim((1,2.5))
  outputfilename="Fig06-Epott_Ekint.png"
  plt.savefig(outputfilename)
  
  fig07 = plt.figure(7)
  plt.rc('text', usetex=True)
  plt.grid()
  plt.xlabel('$ t $')
  plt.ylabel('$ E_{tot} $')
  plt.plot(ttime, Etott,marker='.')
  #plt.plot(ttime, (-0.0615234375)*np.ones(total_time))
  plt.xlim((1,2.5))
  outputfilename="Fig07-Etott.png"
  plt.savefig(outputfilename)
  
  fig08 = plt.figure(8)
  plt.grid()
  plt.xlabel('$ t $')
  plt.ylabel('$ x $')
  plt.plot(ttime, xt[0],marker='.')
  plt.plot(ttime, xt[1],marker='+')
  plt.legend(["$x_0$","$x_1$"])
  outputfilename="Fig08-xt_01.png"
  plt.savefig(outputfilename)
  
  
  fig09 = plt.figure(9)
  plt.grid()
  plt.xlabel('$ t $')
  plt.ylabel('$ vx $')
  plt.plot(ttime, vxt[0],marker='.')
  plt.plot(ttime, vxt[1],marker='+')
  plt.legend(["$v_{x0}$","$v_{x1}$"])
  outputfilename="Fig09-vxt_01.png"
  plt.savefig(outputfilename)
  
  fig10 = plt.figure(10)
  plt.grid()
  plt.xlabel('$ t $')
  plt.ylabel('$ fx $')
  plt.plot(ttime, fxt[0],marker='.')
  plt.plot(ttime, fxt[1],marker='+')
  plt.legend(["$f_{x0}$","$f_{x1}$"])
  outputfilename="Fig10-fxt_01.png"
  plt.savefig(outputfilename)
  #plt.show()
"""     

  #--- Plot   
  plt.figure(3)
  plt.xlabel('$ x $')
  plt.ylabel('$ y $')

  Etott = Ekint+Epott 
  Emax = Etott.max()
  Emin = Etott.min() 
  np.save('Energia.txt', (Ekint, Epott, Etott))
  print '-'*20+'Ekin'
  print Ekint[25:50], len(Ekint), Ekint.max(), Ekint.argmax()
  print '-'*20+'Epot'
  print Etott[25:50], len(Epott), Epott.max(), np.amax(Epott)
  print '-'*20+'Etot'
  print Etott[25:50], len(Etott), Etott.max(), np.amax(Etott)
  # ------ Begin GRAPHICS ------
  plt.grid()  
  plt.figure(1)
  plt.xlabel('$ x $')
  plt.ylabel('$ y $')
  #plt.plot(vxt, vyt)
  plt.xlim((0, gas.N))
  for i in range(gas.N):
    #plt.plot(xt[i,4:9],yt[i,4:9], marker='o')
    plt.plot(xt[i],yt[i])

  plt.figure(2)  
  plt.xlabel('$ t $')
  plt.ylabel('$ E_{tot}$')
  plt.plot(ttime[::ttime_print], Etott, marker='o')
  plt.ylim((Emin-0.00001, Emin+1000))
  
  plt.figure(3)
  plt.xlabel('$ t $')
  plt.ylabel('$ E_{pot}$')
  plt.plot(ttime[::ttime_print], Epott, marker='x')
  
  plt.figure(4)
  plt.xlabel('$ t $')
  plt.ylabel('$ E_{kin}$')
  plt.plot(ttime[::ttime_print], Ekint, marker='+')
  plt.ylim((Ekint.mean()-0.00001, Ekint.mean()+0.00001))
  
  plt.show()

  # Para fazer as figuras 
  gas.arrayPos()
  gas.arrayVel()
  print 'Figure 1'
  # ---- BEGIN FIGURE 1
  plt.figure(1)
  outputfilename = 'LJGas_NpartBoxIniPositions.pdf'
  font = {'family' : 'normal',
        'weight' : 'bold',
        'size'   : 24}
  plt.rc('font', **font)
  plt.rc('text', usetex=True)
  plt.axes().set_aspect('equal')
  plt.yticks(np.arange(0, 2, 0.5))
  plt.xlabel('$ x $')
  plt.ylabel('$ y $')
  area = np.pi*sigma**2
  plt.scatter(gas.X, gas.Y, s=area) # In the scatter the xlim and ylim go after
  plt.xlim((0, gas.xlen))
  plt.ylim((0, gas.ylen))
  
  plt.savefig(outputfilename,transparent=True)
  #plt.show()
  # END FIGURE 1
  
  print 'Figure 2'
  # ---- BEGIN FIGURE 2
  plt.figure(2)
  outputfilename = 'LJGas_NpartBoxIniCond.pdf'
  font = {'family' : 'normal',
        'weight' : 'bold',
        'size'   : 24}
  plt.rc('font', **font)
  plt.rc('text', usetex=True)
  plt.axes().set_aspect('equal')
  plt.yticks(np.arange(0, 2, 0.5))

  ax = plt.gca()
  ax.set_xlabel('$ x $')
  ax.set_ylabel('$ y $')
  ax.quiver(gas.X, gas.Y, gas.VX, gas.VY, color='r')
  plt.scatter(gas.X, gas.Y)
  ax.set_xlim([0,gas.ylen])
  ax.set_ylim([0,gas.ylen])
  plt.savefig(outputfilename,transparent=True)
  #plt.show()
  #plt.close()
  # END FIGURE 2

  gas.arraySpeed()
  import scipy.stats as stats
  maxwell = stats.maxwell
  # data = maxwell.rvs(loc=0, scale=5, size=10000)
  # print data
  params = maxwell.fit(gas.speed, floc=0)

  print 'Figure 3'
  # ---- BEGIN FIGURE 3
  # TODO: Histogram for speed not velocity
  
  plt.figure(3)
  outputfilename = 'LJGas_HistogramSpeed.pdf'
  font = {'family' : 'normal',
        'weight' : 'bold',
        'size'   : 24}
  plt.rc('font', **font)
  plt.rc('text', usetex=True)
  print gas.speed.mean()
  plt.hist(gas.speed, normed=True, bins=20)
  x = np.linspace(0, 15, 100)
  plt.plot(x, maxwell.pdf(x, *params), lw=3)
  print params
  plt.xlabel('$ v $')
  plt.savefig(outputfilename,transparent=True)
  plt.show()
  # END FIGURE 3
  """


